import time
import numpy as np
import faiss


def count_numpy_array(np_arr):  # np_array, n x 1
    key = np.unique(np_arr)
    result = {}
    for k in key:
        mask = (np_arr == k)
        arr_new = np_arr[mask]
        v = arr_new.size
        result[k] = v
    return result


def MAP(preds, all_ids):
    N, d = preds.shape
    # one unique preds_id ===> one unique feature
    preds_ids = np.arange(N)

    preds_ids2_all_ids = dict(zip(preds_ids, all_ids))

    # count true number for specific sku_id
    recall_dict = count_numpy_array(all_ids)

    # index = faiss.index_factory(d,'IDMap,Flat')
    res = faiss.StandardGpuResources()
    # binary
    index = faiss.index_cpu_to_gpu(res, 0, faiss.index_factory(d, 'IDMap,Flat'))

    index.add_with_ids(preds, preds_ids)

    topK = 10

    P_total, R_total = 0.0, 0.0

    batch_search_size = 1024

    for i in range(0, N, batch_search_size):

        # for batch, batch_ids in batch_generator:
        t_start = time.time()
        batch, batch_ids = preds[i:i + batch_search_size, :], all_ids[i:i + batch_search_size]
        _, preds_ids_topK = index.search(batch, topK)

        for search_id, topK_ids in zip(batch_ids, preds_ids_topK):

            True_Count = recall_dict[search_id]
            count, p_tmp, r_tmp = 0.0, 0.0, 0.0
            for ii, preds_id in enumerate(topK_ids):
                if ii == 0: continue  # delete search self
                if search_id == preds_ids2_all_ids[preds_id]:
                    count += 1
                    p_tmp += count / ii

            if count > 0:
                P_total += p_tmp / count
                R_total += count / True_Count
                # print(f'Search_Size:{batch.shape[0]}\t{P_total}\t{R_total}\tSearch Time:{time.time()-t_start}')
        # print(f'Search_Size:{batch.shape[0]}\tSearch Time:{time.time()-t_start}')
        print('\r', f'{i}/{N}', f'Search {preds.shape[0]} Time:{time.time() - t_start}', end='', flush=True)

    return P_total / N, R_total / N
